/*
 * @Description: 该函数不同于seg_flow函数，该函数用来发布分割的真值，用ransac和欧拉聚类来实现分割的真值
 * @Author: Sang Hao
 * @Date: 2021-11-09 19:12:06
 * @LastEditTime: 2021-11-22 21:30:49
 * @LastEditors: Sang Hao
 */

#include <cmath>
#include <pcl/segmentation/extract_clusters.h>
#include "lidar_slam/models/segmentation/ransac_cluster_segment.hpp"

namespace lidar_slam {
/* 首先完成构造函数，对提取平面类做一个初始化 */
RansacClusterSegment::RansacClusterSegment(const YAML::Node& config_node) {
	point_indices_ptr_.reset(new pcl::PointIndices());
	model_coeff_ptr_.reset(new pcl::ModelCoefficients());
	/* 对分割的结果进行初始化 */
	/* 对水平面提取对象初始化 */
	horizontal_plane_extract_ = std::make_shared<PlaneRansacExtract>(
		config_node["horizontal_plane_extract"]);

	/* 对竖直面提取对象初始化 */
	vertical_plane_extract_ = std::make_shared<PlaneRansacExtract>(
		config_node["vertical_plane_extract"]);

	/* 对欧拉聚类参数进行初始化 */
	YAML::Node euclidean_cluster_node =	config_node["euclidean_cluster"];
	min_cluster_size_ = euclidean_cluster_node["min_cluster_size"].as<int>();
	cluster_tolerance_ = euclidean_cluster_node["cluster_tolerance"].as<float>();
}
	
bool RansacClusterSegment::EuclideanClusterExtract(const CloudData::CLOUD_PTR& input_cloud_ptr) {
	/* 这里的输入是去除平面后的点云，用聚类方法实现 */
	/* 创建树，为接下来的聚类做准备 */
	pcl::search::KdTree<CloudData::POINT>::Ptr tree (new pcl::search::KdTree<CloudData::POINT>());
	tree->setInputCloud(input_cloud_ptr);

	/* 接下来做聚类 */
	pcl::EuclideanClusterExtraction<CloudData::POINT> euclidean_cluster_extraction;
	euclidean_cluster_extraction.setClusterTolerance(cluster_tolerance_);
	euclidean_cluster_extraction.setMinClusterSize(min_cluster_size_);
	euclidean_cluster_extraction.setMaxClusterSize(input_cloud_ptr->points.size());
	euclidean_cluster_extraction.setSearchMethod(tree);
	euclidean_cluster_extraction.setInputCloud(input_cloud_ptr);
	euclidean_cluster_extraction.extract(vec_cluster_indices_);

	std::cout << "共计得到cluster数量：" << vec_cluster_indices_.size() << std::endl;

	return true;
}

bool RansacClusterSegment::Segmentation(const CloudData::CLOUD_PTR& input_cloud_ptr, 
		   					   CloudRGBData::CLOUD_PTR& seged_cloud_ptr) {
	float  colors[] = {
		255, 0,   0,   // red 		0	
		0,   255, 0,   // green		1
		0,   0,   255, // blue		2
		255, 255, 0,   // yellow		3
		0,   255, 255, // light blue	4
		255, 0,   255, // magenta     5
		255, 128, 0,   // orange		6
		255, 153, 255, // pink		7
		51,  153, 255, //			8
		153, 102, 51,  //			9
		128, 51,  153, //			10
		153, 153, 51,  //			11
		163, 38,  51,  //			12
		204, 153, 102, //		13
		204, 224, 255, //		14
		128, 179, 255, //		15
		206, 255, 0,   //			16
		255, 204, 204, //			17
		204, 255, 153, //			18
		38,  128, 51,  // 		19

	}; // 20x3=60 color elements
	/**********************************************/
	/* 首先提取地面 */
	/**********************************************/
	std::cout << "输入点云数量：" << input_cloud_ptr->points.size() << std::endl;
	std::cout << "地面点模型为：" << horizontal_plane_extract_->GetModelType() << std::endl;
	std::cout << "开始提取地面..." << std::endl;
	pcl::PointIndices::Ptr point_indices_ptr (new pcl::PointIndices());
	pcl::ModelCoefficients::Ptr model_coeff_ptr (new pcl::ModelCoefficients());
	horizontal_plane_extract_->Extract(input_cloud_ptr, model_coeff_ptr, point_indices_ptr);
	for (auto it = point_indices_ptr->indices.begin(); it != point_indices_ptr->indices.end(); it++) {
		CloudRGBData::POINT p;
		p.x = input_cloud_ptr->points[*it].x;
		p.y = input_cloud_ptr->points[*it].y;
		p.z = input_cloud_ptr->points[*it].z;
		/* 地面点统统设置为黑色 */
		p.r = 0;
		p.g = 0;
		p.b = 0;
		seged_cloud_ptr->points.push_back(p);
	}
	std::cout << "地面点共计 " << point_indices_ptr->indices.size() << "个，已保存" << std::endl;
	std::cout << "地面模型方程：" << model_coeff_ptr->values[0] << "x + "
								<< model_coeff_ptr->values[1] << "y + "
								<< model_coeff_ptr->values[2] << "z + "
								<< model_coeff_ptr->values[3] << " = 0" << std::endl;
	CloudData::CLOUD_PTR temp_cloud_ptr (new CloudData::CLOUD());

	/* 得到分割点云后滤波，用于进一步的分割 */
	pcl::ExtractIndices<CloudData::POINT> horizontal_plane_filter;
	horizontal_plane_filter.setInputCloud(input_cloud_ptr);
	horizontal_plane_filter.setIndices(point_indices_ptr);
	horizontal_plane_filter.setNegative(true);
	horizontal_plane_filter.filter(*temp_cloud_ptr);
	

	/**********************************************/
	/* 接下来分割竖直平面，循环过程中反复对temp_cloud_ptr进行操作 */
	/**********************************************/
	
	std::cout << "竖直平面模型为：" << vertical_plane_extract_->GetModelType() << std::endl;
	for (int i = 0; i < vertical_plane_extract_->GetPlaneNum(); ++i) {
		model_coeff_ptr.reset(new pcl::ModelCoefficients());
		point_indices_ptr.reset(new pcl::PointIndices());
		std::cout << "待平面提取点云尺寸：" << temp_cloud_ptr->points.size() << std::endl;
		vertical_plane_extract_->Extract(temp_cloud_ptr, model_coeff_ptr, point_indices_ptr);

		/* 加入竖直平面点云 */
		for (auto it = point_indices_ptr->indices.begin(); it != point_indices_ptr->indices.end(); it++) {
			CloudRGBData::POINT p;
			p.x = temp_cloud_ptr->points[*it].x;
			p.y = temp_cloud_ptr->points[*it].y;
			p.z = temp_cloud_ptr->points[*it].z;
			p.r = 20 + i * 10;;
			p.g = 20 + i * 10;
			p.b = 20 + i * 10;
			seged_cloud_ptr->points.push_back(p);
		}
		std::cout << "竖直平面 " << i + 1 << "已保存, 共" << point_indices_ptr->indices.size() 
				  << "个点！" << std::endl;
		std::cout << "竖直平面" << i + 1 << " 模型方程：" 
								<< model_coeff_ptr->values[0] << "x + "
								<< model_coeff_ptr->values[1] << "y + "
								<< model_coeff_ptr->values[2] << "x + "
								<< model_coeff_ptr->values[3] << " = 0" << std::endl;  
		
		// CloudData::CLOUD_PTR temp_cloud_ptr (new CloudData::CLOUD());

		/* 得到分割点云后滤波，用于进一步的分割 */
		pcl::ExtractIndices<CloudData::POINT> vertical_plane_filter;
		vertical_plane_filter.setInputCloud(temp_cloud_ptr);
		vertical_plane_filter.setIndices(point_indices_ptr);
		vertical_plane_filter.setNegative(true);
		vertical_plane_filter.filter(*temp_cloud_ptr); /* 滤波后结果依然放在temp_cloud_ptr中 */
	}

	/**********************************************/
	/* 接下来做欧拉聚类，完成障碍物的分割 */
	/**********************************************/

	EuclideanClusterExtract(temp_cloud_ptr);
	for (size_t i = 0; i < vec_cluster_indices_.size(); ++i) {
		for (auto it = vec_cluster_indices_[i].indices.begin(); it != vec_cluster_indices_[i].indices.end(); ++it) {
			CloudRGBData::POINT p;
			p.x = temp_cloud_ptr->points[*it].x;
			p.y = temp_cloud_ptr->points[*it].y;
			p.z = temp_cloud_ptr->points[*it].z;
			p.r = colors[i % 20 * 3 + 0];
			p.g = colors[i % 20 * 3 + 1];
			p.b = colors[i % 20 * 3 + 2];
			seged_cloud_ptr->points.push_back(p);
		}
		std::cout << "非平面点：cluster " << i + 1 << ": 共计" 
				  << vec_cluster_indices_[i].indices.size() << "个点！" 
				  << std::endl;
	}
	seged_cloud_ptr->height = 1;
	seged_cloud_ptr->width = seged_cloud_ptr->points.size();
	return true;
}
}
